from functools import lru_cache
import torch
from torch import nn


"""
这段代码实现了 RoPE（Rotary Position Embedding，旋转位置编码） 的核心逻辑，这是一种在大语言模型（如 LLaMA、GPT-4）中广泛使用的位置编码技术，通过对注意力的 Query（Q）和 Key（K）向量进行旋转变换，注入位置信息，同时保持注意力的相对位置不变性。以下从「核心原理→代码拆解→优化细节」展开分析：
一、核心背景：RoPE 是什么？
传统位置编码（如 Transformer 原论文的正弦编码）是将位置信息 “叠加” 到词嵌入中，而 RoPE 则通过旋转矩阵直接对 Q/K 向量进行变换，核心优势是：

相对位置不变性：Q 与 K 的注意力得分仅依赖两者的相对位置，而非绝对位置（符合语言的上下文关联性）。
长序列泛化：支持超过训练时最大长度的序列（通过外推）。
计算高效：可在注意力计算前对 Q/K 批量处理，无额外注意力层开销。

RoPE 的核心公式（简化版）：
对 Q/K 向量的每对维度（如第 2i 和 2i+1 维）应用旋转矩阵

二、代码拆解：从工具函数到类实现
代码分为 3 部分：apply_rotary_emb（旋转变换工具函数）、RotaryEmbedding（RoPE 核心类）、get_rope（单例工厂函数）。
三、关键优化点
预计算缓存：
初始化时计算所有位置的 cos/sin 并缓存，推理时直接索引读取，避免每批 token 重复计算三角函数（三角函数计算开销较高，缓存后可提升 10~100 倍速度）。
torch.compile 加速：
forward 方法用 @torch.compile 装饰，PyTorch 会自动优化计算图（如算子融合、内存复用），尤其在批量处理大量 token 时，可显著降低 Python 层 overhead。
广播机制：
cos/sin.unsqueeze(-2) 扩展维度后，通过广播适配多头结构，无需循环处理每个注意力头，提升并行效率。
非持久化缓存：
register_buffer(..., persistent=False) 确保缓存不保存到模型 checkpoint，避免 checkpoint 体积过大（缓存仅在内存中使用，无需持久化）。
四、应用场景与使用流程
RoPE 是 LLM 注意力层的前置步骤，典型使用流程如下：

初始化 RoPE：通过 get_rope 创建单例实例（指定 head_size、max_position 等参数）。
获取位置列表：对输入序列，生成每个 token 的位置（如 positions = torch.arange(seq_len, device=device)）。
旋转 Q/K：在注意力计算前，调用 rotary_emb.forward(positions, query, key)，得到旋转后的 Q/K。
计算注意力：用旋转后的 Q/K 进行注意力计算（如 FlashAttention），此时注意力得分已包含位置信息。
总结
这段代码是工业级 LLM 中 RoPE 的标准实现，核心特点是：

正确性：严格遵循 RoPE 旋转矩阵公式，确保位置信息准确注入。
高效性：预计算缓存 + torch.compile 优化，批量处理速度快。
易用性：单例工厂函数 + 批量接口，适配多轮推理和长序列场景。

理解此代码后，可清晰掌握 RoPE 从 “数学公式” 到 “工程实现” 的转化逻辑，以及如何通过缓存、广播等优化手段平衡性能与内存开销。
"""


def apply_rotary_emb(
        x: torch.Tensor,  # 输入Q/K向量，形状通常为 [num_tokens, num_heads, head_size]
        cos: torch.Tensor,  # 余弦值缓存，形状 [num_tokens, rotary_dim//2]
        sin: torch.Tensor,  # 正弦值缓存，形状 [num_tokens, rotary_dim//2]
) -> torch.Tensor:
    """
    执行旋转变换
关键细节：
数据类型转换：先将 x 转为 float32 计算（低精度如 float16 可能导致旋转后数值不稳定），最后恢复原始类型。
维度匹配：cos/sin.unsqueeze(-2) 增加 num_heads 维度（1），通过广播机制适配 Q/K 的多头结构（无需手动循环处理每个头）。
分块旋转：按最后一维拆分为两部分，直接应用旋转矩阵，避免逐维度循环，提升计算效率。
    """
    # 1. 扩展cos/sin维度，适配Q/K的多头结构（增加num_heads维度）
    cos = cos.unsqueeze(-2)  # 形状：[num_tokens, 1, rotary_dim//2]
    sin = sin.unsqueeze(-2)  # 形状：[num_tokens, 1, rotary_dim//2]

    # 2. 将x按最后一维（head_size）拆分为两部分（x1对应偶数维，x2对应奇数维）
    x1, x2 = torch.chunk(x.to(torch.float32), 2, dim=-1)  # 各部分形状：[num_tokens, num_heads, rotary_dim//2]

    # 3. 应用旋转矩阵公式（对应上文核心公式）
    y1 = x1 * cos - x2 * sin  # 旋转后的偶数维
    y2 = x2 * cos + x1 * sin  # 旋转后的奇数维

    # 4. 拼接结果并恢复原始数据类型（避免低精度计算损失）
    return torch.cat((y1, y2), dim=-1).to(x.dtype)


class RotaryEmbedding(nn.Module):
    def __init__(
            self,
            head_size: int,  # 每个注意力头的维度（如 128）
            rotary_dim: int,  # RoPE 作用的维度（通常等于 head_size，即对所有维度应用旋转）
            max_position_embeddings: int,  # 预计算的最大位置（如 2048、4096）
            base: float,  # 衰减基数（通常为 10000，控制位置衰减速度）
    ) -> None:
        """
        RoPE 核心类（预计算缓存 + 批量处理）
        核心逻辑拆解：
预计算缓存：初始化时计算所有位置（0~max_position-1）的 cos 和 sin，存储为 cos_sin_cache，避免推理时重复计算（大幅提升速度）。
批量位置处理：forward 接收 positions（如当前批次的 token 位置），直接从缓存中读取对应位置的 cos/sin，支持任意位置的批量处理（如非连续位置、长序列切片）。
形状适配：通过 view 重塑 Q/K 为标准形状（[num_tokens, num_heads, head_size]），确保旋转函数可正确处理，最后恢复原始形状（适配上层输入格式）。
        """
        super().__init__()
        self.head_size = head_size
        assert rotary_dim == head_size  # 确保 RoPE 作用于所有头维度（简化实现）

        # -------------------------- 步骤1：计算频率矩阵（inv_freq + freqs） --------------------------
        # inv_freq：每个维度对的衰减频率，形状 [rotary_dim//2]
        # 公式：inv_freq[i] = 1 / (base^(2i/rotary_dim)) （i 是维度对索引，0~rotary_dim//2-1）
        inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))

        # t：位置序列，形状 [max_position_embeddings]（0 ~ max_position_embeddings-1）
        t = torch.arange(max_position_embeddings, dtype=torch.float)

        # freqs：位置-维度对的频率矩阵，形状 [max_position_embeddings, rotary_dim//2]
        # 公式：freqs[p][i] = p * inv_freq[i] （p 是位置，i 是维度对索引）
        freqs = torch.einsum("i,j -> ij", t, inv_freq)  # 矩阵外积：t（1D）× inv_freq（1D）→ 2D

        # -------------------------- 步骤2：预计算 cos/sin 缓存 --------------------------
        # 对每个位置-维度对计算 cos 和 sin，拼接后作为缓存
        cos = freqs.cos()  # 形状 [max_position_embeddings, rotary_dim//2]
        sin = freqs.sin()  # 形状 [max_position_embeddings, rotary_dim//2]
        cache = torch.cat((cos, sin), dim=-1)  # 拼接后：[max_position_embeddings, rotary_dim]

        # -------------------------- 步骤3：注册缓存（非持久化，节省内存） --------------------------
        # register_buffer：将缓存注册为模型缓冲区（非可学习参数，随模型移动设备）
        # persistent=False：训练时不保存到 checkpoint（避免缓存占用过多空间）
        self.register_buffer("cos_sin_cache", cache, persistent=False)

    @torch.compile  # 启用 PyTorch 编译优化，提升批量处理速度
    def forward(
            self,
            positions: torch.Tensor,  # 需处理的位置列表，形状 [num_tokens]（如 [0,1,2,3] 表示前4个位置）
            query: torch.Tensor,  # 输入 Query 向量，形状 [num_tokens, num_heads, head_size]
            key: torch.Tensor,  # 输入 Key 向量，形状 [num_tokens, num_heads, head_size]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # 1. 获取当前批次的 token 数量
        num_tokens = positions.size(0)

        # 2. 从预计算缓存中读取对应位置的 cos/sin（避免重复计算）
        cos_sin = self.cos_sin_cache[positions]  # 形状 [num_tokens, rotary_dim]
        cos, sin = cos_sin.chunk(2, dim=-1)  # 拆分后各为 [num_tokens, rotary_dim//2]

        # 3. 处理 Query：重塑形状→旋转→恢复形状（适配批量处理）
        query_shape = query.shape  # 保存原始形状（避免后续维度不匹配）
        # 重塑为 [num_tokens, num_heads, head_size]（确保最后一维是 head_size）
        query = query.view(num_tokens, -1, self.head_size)
        query = apply_rotary_emb(query, cos, sin).view(query_shape)  # 旋转后恢复原始形状

        # 4. 处理 Key：与 Query 逻辑一致
        key_shape = key.shape
        key = key.view(num_tokens, -1, self.head_size)
        key = apply_rotary_emb(key, cos, sin).view(key_shape)

        # 5. 返回旋转后的 Q 和 K
        return query, key


@lru_cache(1)  # 装饰器：缓存结果，最多缓存1个实例（单例模式）
def get_rope(
    head_size: int,
    rotary_dim: int,
    max_position: int,
    base: float,
    rope_scaling: dict | None = None,  # 预留缩放参数（当前未实现）
):
    """
    单例工厂函数（避免重复创建）
核心作用：通过 lru_cache(1) 确保全局只创建一个 RotaryEmbedding 实例，避免重复初始化缓存（节省内存，尤其在多轮推理中）。
预留扩展：rope_scaling 参数用于后续支持 “RoPE 缩放”（如通过插值适配更长序列），当前简化实现中未启用。
    """
    assert rope_scaling is None  # 暂不支持 RoPE 缩放（如动态缩放长序列）
    # 创建并返回 RotaryEmbedding 实例（缓存后重复调用返回同一实例）
    rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base)
    return rotary_emb
